# -*- coding: utf-8 -*-
"""
Created on Thu Jul 25 19:31:03 2024

@author: xiaok
"""
import mne
import os
import json

#%%
# use multiple files in a folder by list
# abandon using concatenate_raws

trianingDataFolder = 'D:\\m_proj_24\\mii_app\\training_data\\'

# file_list=[]
raw_list = []

for file_name in os.listdir(trianingDataFolder):
    filePathName = trianingDataFolder+file_name
    # file_list.append(filePathName)
    if not (filePathName.endswith('_.edf')):
        continue
    
    raw = mne.io.read_raw_edf(filePathName,preload=True) # preload should be True, if you want to use filter like below 

    raw.filter(l_freq=1, h_freq=90, method='iir')
    raw.notch_filter(freqs=50)

    raw_list.append(raw)

final_raw = mne.concatenate_raws(raw_list)


#%%
with open('D:\\m_proj_24\\mii_app\\task\\task_markers.json', 'r') as file:
    markers = json.load(file)
    
label_cue_left_int = markers['Cue_onset_left'][0]
label_cue_right_int = markers['Cue_onset_right'][0]
label_cue_up_int = markers['Cue_onset_up'][0]
label_cue_down_int = markers['Cue_onset_down'][0]    

events_from_annot,event_dict = mne.events_from_annotations(final_raw)
# List of the events
# "trial_end":[13],             1
# "Cue_onset_left":[22],        2
# "Cue_onset_right":[23],       3
# "Cue_onset_up":[24],          4
# "Cue_onset_down":[25],        5
# "begin":[99],                 6

# fig = mne.viz.plot_events(events_from_annot,event_id=event_dict,sfreq=raw.info['sfreq'],first_samp=raw.first_samp)

event_dict_new = {value:int(float(key)) for key,value in event_dict.items()}
events_from_annot_new = events_from_annot.copy()
for i, c in enumerate(events_from_annot[:,2]):
    events_from_annot_new[i,2]=event_dict_new[c]

evnet_dict_marker = {
    'left':label_cue_left_int,
    'right':label_cue_right_int,
    'up':label_cue_up_int,
    'down':label_cue_down_int
}

#%%

signal_win_start = 0.5 #second
signal_win_end = 2.5 #second

epochs = mne.Epochs(
    final_raw,
    events_from_annot_new,
    event_id=evnet_dict_marker,
    tmin=signal_win_start,
    tmax=signal_win_end,
    baseline = None
)

data = epochs.get_data()
labels = epochs.events[:,-1]


#%%

#%%
save_pickle = 0
import pickle

mMIIData = {}
mMIIData['datas']=data
mMIIData['labels']=labels
mMIIData['fs']=raw.info['sfreq']

filePathName_save = 'D:\\m_proj_24\\mii_app\\training_data\\mMIIData.p'

if save_pickle==1:
    pickle.dump(mMIIData,open(filePathName_save,'wb'))

#%%
#%%

# raw_list = []
# evt_list=[]
# evt_dic_list = []
# for f in file_list:
#     raw_tmp = mne.io.read_raw_edf(f,preload=False)
#     raw_list.append(raw_tmp)
#     events_from_annot,event_dict = mne.events_from_annotations(raw_tmp)
#     evt_list.append(events_from_annot)
#     evt_dic_list.append(event_dict)


# raw_=mne.io.concatenate_raws(raw_list,events_list=evt_list)
# raw = raw_[0]
# evt = raw_[1]

# fs = raw.info['sfreq']
# nchannels, nsamples = raw.get_data().shape      


    
#%%
# with open('D:\\m_proj_24\\mii_app\\task\\task_markers.json', 'r') as file:
#     markers = json.load(file)
    
# label_cue_left_int = markers['Cue_onset_left'][0]
# label_cue_right_int = markers['Cue_onset_right'][0]
# label_cue_up_int = markers['Cue_onset_up'][0]
# label_cue_down_int = markers['Cue_onset_down'][0]    

# events_from_annot,event_dict = mne.events_from_annotations(raw)
# # List of the events
# # "trial_end":[13],             1
# # "Cue_onset_left":[22],        2
# # "Cue_onset_right":[23],       3
# # "Cue_onset_up":[24],          4
# # "Cue_onset_down":[25],        5
# # "begin":[99],                 6

# # fig = mne.viz.plot_events(events_from_annot,event_id=event_dict,sfreq=raw.info['sfreq'],first_samp=raw.first_samp)

# event_dict_new = {value:int(float(key)) for key,value in event_dict.items()}
# events_from_annot_new = events_from_annot.copy()
# for i, c in enumerate(events_from_annot[:,2]):
#     events_from_annot_new[i,2]=event_dict_new[c]

# evnet_dict_marker = {
#     'left':label_cue_left_int,
#     'right':label_cue_right_int,
#     'up':label_cue_up_int,
#     'down':label_cue_down_int
# }

# signal_win_start = 0.5 #second
# signal_win_end = 2.5 #second

# epochs = mne.Epochs(
#     raw,
#     events_from_annot_new,
#     event_id=evnet_dict_marker,
#     tmin=signal_win_start,
#     tmax=signal_win_end,
#     baseline = None
# )

# data = epochs.get_data()
# labels = epochs.events[:,-1]

#%%
#%%
#%%
#%%
#%%
    
    
